
#include "ICMPDef.h"
#include "util.c"


struct node* create_and_add_node(struct node *path, char *ip){
	struct node *aux = malloc(sizeof(struct node));
	aux->addr = ip;
	aux->next = NULL;

	//path = add_node(path, aux);

	if (path == NULL) {
		path = aux;
	} else {
		struct node *tmp = path;
		while (tmp->next != NULL) {
			tmp = tmp->next;
		}
		tmp->next = aux;
	}

	return path;
}

unsigned short check_sum(unsigned short *buf, int size){
    unsigned long sum=0;
    int i;
    for (i = 0; i < size; i++){
        sum += *buf++;
    }
    sum = (sum >> 16) + (sum & 0xffff);
    sum += (sum >> 16);
    return ~sum;
}

int get_socket() {
	int fd;
	int one = 1;
	struct timeval Timeout;

	fd = socket (PF_INET, SOCK_RAW, IPPROTO_ICMP);

	if (fd < 0) {
		printf ("ERROR: Could not open socket.\n");
	}

	// Seta valores de Timeout
	Timeout.tv_sec = SOCKET_TIMEOUT;
	Timeout.tv_usec = 0;

	if (setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) != 0 ) {
		printf("ERROR: Could not set SO_REUSEADDR\n");
	}
	if (setsockopt(fd, SOL_IP, IP_TTL, &one, sizeof(one)) != 0 ) {
		printf("ERROR: Could not set IP_TTL\n");
	}
	if (setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &Timeout, sizeof(Timeout)) != 0 ) {
		printf ("ERROR: Could not set SO_RCVTIMEO\n");
	}

	return fd;
}

void set_socket_ttl(int sd, int ttl) {
	if (setsockopt(sd, SOL_IP, IP_TTL, &ttl, sizeof(ttl)) != 0 ) {
		printf("ERROR: Could not set IP_TTL\n");
	}
}

char* create_packet() {
	int i;
	//char buffer[PACKET_SIZE];
	char *buffer = malloc(PACKET_SIZE * sizeof(char));
	struct icmphdr *icmp_header = (struct icmphdr *) buffer;

	for (i=0; i<= PACKET_SIZE; i++) {
		buffer[i] = 0;
	}

	// Preenchimento do cabecalho ICMP
    icmp_header->type = ICMP_ECHO;
    icmp_header->code = 0;
    icmp_header->un.echo.id = htons(0);
    icmp_header->un.echo.sequence = htons(0);
    icmp_header->checksum = check_sum((unsigned short *)icmp_header, ICMP_SIZE);

	return buffer;
}

char* receive_packet(int sd) {
	struct sockaddr_in addr;
	int addr_size = sizeof(addr);
	char *buffer = malloc(PACKET_SIZE * sizeof(char));
	char *ip_addr;

	int recv_b = recvfrom(sd,buffer,PACKET_SIZE,MSG_WAITALL,(struct sockaddr*)&addr,&addr_size);
	if (recv_b < 0) {
		return "*";
	}

	ip_addr = (char *) malloc(16);
	strcpy(ip_addr, (char *) inet_ntoa( addr.sin_addr.s_addr));
	return ip_addr;
}

struct node* trace_route(struct in_addr dest_addr) {
	int has_finished = FALSE;
	int step = 0;
	struct node *path = NULL;
	struct sockaddr_in dest;

	char* dest_ip_addr = (char *) malloc(16);
	strcpy(dest_ip_addr, (char *) inet_ntoa( dest_addr.s_addr));

	int sd = get_socket();
	if (sd  < 0) {
		return (void*)-1;
	}

	char *packet_snd = create_packet();

	dest.sin_family = AF_INET;
	dest.sin_port = 0;
	dest.sin_addr.s_addr = dest_addr.s_addr;

	while (!has_finished && (step < MAX_STEPS)) {
		step++;
		set_socket_ttl(sd, step);

		int s = sendto (sd,packet_snd,PACKET_SIZE,0,(struct sockaddr *) &dest,sizeof (dest));
		if (s < 0){
			printf("ERROR: Could not send message %d, %d\n", sd, errno);
		}
	
		char *rcv_ip_addr = receive_packet(sd);
		printf("Response from %s\n", rcv_ip_addr);

		if (!strcmp(dest_ip_addr, rcv_ip_addr)) {
			has_finished = TRUE;
		}
		path = create_and_add_node(path, rcv_ip_addr);
	}

	return path;
}

void receive_msg(int con_fd) {
	ssize_t n;
	struct in_addr dest_addr;
	struct probe_args args;
	struct probe_result result;

	for ( ; ; ) {
		if ( (n = readn(con_fd, &args, sizeof(args))) == 0) {
			return;   // connection closed
		}
		//bzero(&args, sizeof(args));
		bzero(&result, sizeof(result));

		printf("probe.c: received request to discover path to %s\n", args.addr);

		dest_addr.s_addr = inet_addr(args.addr);
		struct node* path = trace_route(dest_addr);
	
		if ((int)path == -1) {
			printf("ERROR: Could not trace the route to destination.\n");
			result.error = TRUE;
		} else if (path == NULL) {
			printf("No path returned\n");
		} else {
			struct node* n = path;
			int hosts = 0;
			while (n != NULL) {
				//printf("%d: %s\n", hosts, n->addr);
				strncpy(result.res[hosts].ip_addr, n->addr, 16);
				result.res[hosts].ip_addr[16] = '\0';
				n = n->next;
				hosts++;
			}
			result.num_hosts = hosts + 1;
			printf("probe.c: sending response\n");
			writen(con_fd, &result, sizeof(result));
		}
	}

	dest_addr.s_addr = inet_addr("209.85.193.99");

	struct node* path = trace_route(dest_addr);

	if ((int)path == -1) {
		printf("ERROR: Could not trace the route to destination.\n");
	} else if (path == NULL) {
		printf("No path returned\n");
	} else {
		struct node* n = path;
		int hosts = 1;
		while (n != NULL) {
			//printf("%d: %s\n", hosts, n->addr);
			n = n->next;
			hosts++;
		}
	}
}

int start_listener(int bind_port) {
	int list_fd;
	struct sockaddr_in servaddr;

	list_fd = socket (AF_INET, SOCK_STREAM, 0);

	bzero(&servaddr, sizeof(servaddr));
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = htonl (INADDR_ANY);
	servaddr.sin_port = htons (bind_port);

	bind(list_fd, (struct sockaddr *) &servaddr, sizeof(servaddr));

	listen(list_fd, PROBE_MAX_CON);
	
	return list_fd;
}

void probe_loop(int list_fd) {
	int con_fd;
	pid_t child_pid;
	struct sockaddr_in client_addr;
	socklen_t client_len = sizeof(client_addr);

	for ( ; ; )  {
		con_fd = accept(list_fd, (struct sockaddr *) &client_addr, &client_len);

		if ( (child_pid = fork()) == 0) { /* child process */
			close(list_fd);    /* close listening socket */
			receive_msg(con_fd);   /* process the request */
			exit (0);
		}
		close(con_fd);          /* parent closes connected socket */
	}
}

// ----------------------------------------------------------------------------
// Principal-------------------------------------------------------------------
int main(int argc, char *argv[]) {

	if (argc < 2) {
		printf("Usage: probe <port>\n");
		return -1;
	}

	int bind_port = atoi(argv[1]);
	printf("probe.c: Binding to port %d\n", bind_port);
	int list_fd = start_listener(bind_port);	

	/* start main loop   */
	probe_loop(list_fd);

/*
	struct in_addr dest_addr;

	dest_addr.s_addr = inet_addr("209.85.193.99");

	struct node* path = trace_route(dest_addr);

	if ((int)path == -1) {
		printf("ERROR: Could not trace the route to destination.\n");
	} else if (path == NULL) {
		printf("No path returned\n");
	} else {
		struct node* n = path;
		int hosts = 1;
		while (n != NULL) {
			printf("%d: %s\n", hosts, n->addr);
			n = n->next;
			hosts++;
		}
	}
*/

    return 0;
}
